package cert_utils

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"io/ioutil"
	"log"
	"math/big"
	"net"
	"os"
	"time"
)

func GenerateX509Cert(commonName string, countries, hosts, orgs []string,
	notBefore, notAfter time.Time, keyBits int,
	parentCert *x509.Certificate, parentPrivKey *rsa.PrivateKey) (key *rsa.PrivateKey, cert *x509.Certificate) {
	key, err := rsa.GenerateKey(rand.Reader, keyBits)
	if err != nil {
		log.Fatalf("Failed to generate private key: %s!\n", err)
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		log.Fatalf("Failed to generate serial number: %s!\n", err)
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   commonName,
			Country:      countries,
			Organization: orgs,
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
		IsCA: false,
	}

	for _, host := range hosts {
		if ip := net.ParseIP(host); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, host)
		}
	}

	if parentCert == nil {
		template.IsCA = true
		template.KeyUsage |= x509.KeyUsageCertSign
		parentCert = &template
		parentPrivKey = key
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, parentCert, &key.PublicKey, parentPrivKey)
	cert, err = x509.ParseCertificate(derBytes)
	if err != nil {
		log.Fatalf("Failed to parse generated cert: %s!\n", err)
	}
	return
}

func ReadCert(fileName string) *x509.Certificate {
	certBytes, err := ioutil.ReadFile(fileName)
	if err != nil {
		log.Fatalf("Failed to read %s: %s\n", fileName, err)
	}
	block, rest := pem.Decode(certBytes)
	if block == nil || len(rest) > 0 {
		log.Fatalf("Invalid pem file: %s\n", fileName)
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		log.Fatalf("Invalid certificate %s: %s\n", fileName, err)
	}
	return cert
}

func WriteCert(cert *x509.Certificate, fileName string) {
	certOutFile, err := os.Create(fileName)
	defer certOutFile.Close()
	if err != nil {
		log.Fatalf("Failed to open %s for writing: %s!\n", fileName, err)
	}
	pem.Encode(certOutFile, &pem.Block{Type: "CERTIFICATE", Bytes: cert.Raw})
}

func ReadKey(fileName string) *rsa.PrivateKey {
	keyBytes, err := ioutil.ReadFile(fileName)
	if err != nil {
		log.Fatalf("Failed to read %s: %s\n", fileName, err)
	}
	block, rest := pem.Decode(keyBytes)
	if block == nil || len(rest) > 0 {
		log.Fatalf("Invalid pem file: %s\n", fileName)
	}
	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		log.Fatalf("Invalid private key %s: %s\n", fileName, err)
	}
	return key
}

func WriteKey(key *rsa.PrivateKey, fileName string) {
	keyOutFile, err := os.Create(fileName)
	defer keyOutFile.Close()
	if err != nil {
		log.Fatalf("Failed to open %s for writing: %s!\n", fileName, err)
	}
	pem.Encode(keyOutFile, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
}
